#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable

#define ACCUM_FLT4 float4
#define FLT float
#define FLT2 float2
#define FLT3 float3
#define FLT4 float4
#define TO_FLT4 convert_float4
#define TO_ACCUM_TYPE convert_float4
#define TO_ACCUM_FLT convert_float
#define READ_IMAGE read_imagef
#define WRITE_IMAGE write_imagef

__constant sampler_t smp_edge = CLK_NORMALIZED_COORDS_FALSE |
                                CLK_ADDRESS_CLAMP_TO_EDGE | CLK_FILTER_NEAREST;

__constant sampler_t smp_none =
    CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

__constant sampler_t smp_zero =
    CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;

// num_group[0]=(1,112,1)
// local_size[0]=(56,1,8)
// global_size[0]=(112,112,8)
__kernel void
main_function(__read_only image2d_t src_data,  // WHC?
              __global float4 *filters_buffer, // ?
              __global float4 *biases,
              float relu_clip1, // 6.0
              __write_only image2d_t dst_data,
              int4 stride_padding,  // 2 2 0 0
              int4 kernel_dilation, // 3 3 1 1
              // int Slices() const { return DivideRoundUp(shape_.c, 4); }
              // int4(shape_.w * shape_.b, shape_.h, Slices(), shape_.b);
              int4 src_size, // 224 224 1 1
              int4 dst_size) // 112 112 8 1
{

  // conv_params.block_size[0] =  2
  // conv_params.block_size[1] =  1
  // conv_params.block_size[2] =  2
  int X = get_global_id(0) * 2; // 0-224
  int Y = get_global_id(1) * 1; // 0-112
  int Z = get_global_id(2) * 2; // 0-16
  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) {
    return;
  }
  // X = 0-112
  // Y = 0-112
  // Z = 8

  ACCUM_FLT4 r000 = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
  ACCUM_FLT4 r001 = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
  ACCUM_FLT4 r100 = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
  ACCUM_FLT4 r101 = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);

  // 在输入上的偏移
  // xc0 = (X + 0) * stride[0] - pad[0] = (0-112) * 2
  int xc0 = (X + 0) * stride_padding.x + stride_padding.z;
  // xc1 = (X + 1) * stride[0] - pad[0] = (1-113) * 2
  int xc1 = (X + 1) * stride_padding.x + stride_padding.z;
  // yc0 = (Y + 0) * strid3[1] - pad[1] = (0-112) * 2
  int yc0 = (Y + 0) * stride_padding.y + stride_padding.w;

  __global float4 *weights_cache;
  __global float4 *filters_loc = filters_buffer + Z * 4 * src_size.z *
                                                      kernel_dilation.x *
                                                      kernel_dilation.y;

  for (int ky = 0; ky < kernel_dilation.y; ++ky) { // 0~kernel[1]  0~3
    // kernel_dilation.w = dilation[1] = 1
    int yck0 = ky * kernel_dilation.w + yc0;         // yc0+(0-3)
    for (int kx = 0; kx < kernel_dilation.x; ++kx) { // 0~kernel[0]  0~3
      // kernel_dilation.z = dilation[0] = 1
      int xck0 = kx * kernel_dilation.z + xc0;
      int xck1 = kx * kernel_dilation.z + xc1;
      int s = 0;
      do {
        float4 src00;
        float4 src01;
        weights_cache = filters_loc;
        src00 = read_imagef(src_data, smp_zero,
                            (int2)((xck0), (yck0)*src_size.z + (s)));
        src01 = read_imagef(src_data, smp_zero,
                            (int2)((xck1), (yck0)*src_size.z + (s)));
        s += 1;
        r000 += weights_cache[0] * src00.x;
        r001 += weights_cache[0] * src01.x;
        r000 += weights_cache[1] * src00.y;
        r001 += weights_cache[1] * src01.y;
        r000 += weights_cache[2] * src00.z;
        r001 += weights_cache[2] * src01.z;
        r000 += weights_cache[3] * src00.w;
        r001 += weights_cache[3] * src01.w;
        r100 += weights_cache[4] * src00.x;
        r101 += weights_cache[4] * src01.x;
        r100 += weights_cache[5] * src00.y;
        r101 += weights_cache[5] * src01.y;
        r100 += weights_cache[6] * src00.z;
        r101 += weights_cache[6] * src01.z;
        r100 += weights_cache[7] * src00.w;
        r101 += weights_cache[7] * src01.w;
        filters_loc += 8;
      } while (s < src_size.z);
    };
  };

  weights_cache = biases + Z;

  if (Z + 0 >= dst_size.z)
    return;
  {
    FLT4 bias_val = TO_FLT4(weights_cache[0]);
    {
      FLT4 res = TO_FLT4(r000) + bias_val;
      { res = clamp(res, (FLT)(0.0f), (FLT)(relu_clip1)); }
      write_imagef(dst_data,
                   (int2)(((X + 0)), ((Y + 0)) * dst_size.z + (Z + 0)), res);
    }
    if ((X + 1) < dst_size.x) {
      FLT4 res = TO_FLT4(r001) + bias_val;
      { res = clamp(res, (FLT)(0.0f), (FLT)(relu_clip1)); }
      write_imagef(dst_data,
                   (int2)(((X + 1)), ((Y + 0)) * dst_size.z + (Z + 0)), res);
    }
  }

  if (Z + 1 >= dst_size.z)
    return;

  {
    FLT4 bias_val = TO_FLT4(weights_cache[1]);
    {
      FLT4 res = TO_FLT4(r100) + bias_val;
      { res = clamp(res, (FLT)(0.0f), (FLT)(relu_clip1)); }
      write_imagef(dst_data,
                   (int2)(((X + 0)), ((Y + 0)) * dst_size.z + (Z + 1)), res);
    }
    if ((X + 1) < dst_size.x) {
      FLT4 res = TO_FLT4(r101) + bias_val;
      { res = clamp(res, (FLT)(0.0f), (FLT)(relu_clip1)); }
      write_imagef(dst_data,
                   (int2)(((X + 1)), ((Y + 0)) * dst_size.z + (Z + 1)), res);
    }
  }
}

#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
#define ACCUM_FLT4 float4
#define FLT float
#define FLT2 float2
#define FLT3 float3
#define FLT4 float4
#define TO_FLT4 convert_float4
#define TO_ACCUM_TYPE convert_float4
#define TO_ACCUM_FLT convert_float
#define READ_IMAGE read_imagef
#define WRITE_IMAGE write_imagef
__constant sampler_t smp_edge = CLK_NORMALIZED_COORDS_FALSE |
                                CLK_ADDRESS_CLAMP_TO_EDGE | CLK_FILTER_NEAREST;
__constant sampler_t smp_none =
    CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;
__constant sampler_t smp_zero =
    CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;

__kernel void main_function(__read_only image2d_t src_data,
                            __global float4 *filters_buffer,
                            __global float4 *biases, float relu_clip1,
                            __write_only image2d_t dst_data,
                            int4 stride_padding, int4 kernel_dilation,
                            int4 src_size, int4 dst_size) {
  int X = get_global_id(0) * 1;
  int Y = get_global_id(1) * 1;
  int Z = get_global_id(2) * 1;
  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) {
    return;
  }
  ACCUM_FLT4 r000 = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
  int xc0 = (X + 0) * stride_padding.x + stride_padding.z;
  int yc0 = (Y + 0) * stride_padding.y + stride_padding.w;
  __global float4 *weights_cache;
  __global float4 *filters_loc = filters_buffer + Z * 4 * src_size.z *
                                                      kernel_dilation.x *
                                                      kernel_dilation.y;
  for (int ky = 0; ky < 3; ++ky) {
    int yck0 = ky + yc0;
    for (int kx = 0; kx < 3; ++kx) {
      int xck0 = kx + xc0;

      float4 src00;
      weights_cache = filters_loc;
      src00 = read_imagef(src_data, smp_zero, (int2)(xck0, yck0 * src_size.z));
      r000 += weights_cache[0] * src00.x;
      r000 += weights_cache[1] * src00.y;
      r000 += weights_cache[2] * src00.z;
      r000 += weights_cache[3] * src00.w;
      filters_loc += 4;
    };
  };

  weights_cache = biases + Z;
  if (Z + 0 >= dst_size.z)
    return;
  {
    FLT4 bias_val = TO_FLT4(weights_cache[0]);
    {
      FLT4 res = TO_FLT4(r000) + bias_val;
      { res = clamp(res, (FLT)(0.0f), (FLT)(relu_clip1)); }
      write_imagef(dst_data, (int2)(X, Y * dst_size.z + Z), res);
    }
  }
}

#pragma OPENCL EXTENSION cl_khr_fp16 : enable

// grid[0]=224 grid[1]=224 grid[2]=1
// work_group_size=local
// local[0]=16 local[1]=8  local[2]=1
// global[0]=224 global[1]=224 global[2]=1

__kernel void to_tensor(__global float *src, __write_only image2d_t dst,
                        int4 size, int channels) {
  int linear_id = get_global_id(0);
  int x = (linear_id / size.w);
  int b = linear_id % size.w;
  int y = get_global_id(1);
  int d = get_global_id(2);

  if (x >= size.x || y >= size.y || d >= size.z)
    return;

  float4 result;
  int c = d * 4;
  int index = ((b * size.y + y) * size.x + x) * channels + c;
  result.x = src[index];
  result.y = c + 1 < channels ? src[index + 1] : 1;
  result.z = c + 2 < channels ? src[index + 2] : 2;
  result.w = c + 3 < channels ? src[index + 3] : 3;

  write_imagef(dst, (int2)((x)*size.w + (b), (y)*size.z + (d)), result);
}